#include <stdio.h>
#include <stdlib.h>
#include <math.h>

__global__ void jacobi(float* AD, float* XD, float* XkD, int* iteracaoD, int ordem, float erro, int MaxIteracao, int n_threads){

    int tid = threadIdx.x;
    int i, iteracao=0;
    __shared__ float auxErro, auxErro1;

    if(tid < ordem){
        while(iteracao < MaxIteracao){    //Repetiçao ateh o numero  de iteraçoes que foi passada pelo usuario.
            iteracao++;
            while(tid < ordem){

                XkD[tid]=0;

                for(i=0;i<ordem;i++){
                    if(tid == i)
                        XkD[tid] += AD[tid*ordem+tid];
                    else
                        XkD[tid] += AD[tid*ordem+i]*XD[i];
                }
                tid+=n_threads;
            }

            auxErro = 0;
            auxErro1 = 0;
            __syncthreads();

            tid = threadIdx.x;

            while(tid < ordem){
                if(fabs(XkD[tid] - XD[tid]) > auxErro)   //Parte em que se calcula o modulo dos valores, para calcular o erro.
                    auxErro = fabs(XkD[tid] - XD[tid]);

                if(fabs(XkD[tid]) > auxErro1)   //Parte em que se calcula o modulo dos valores, para calcular o erro.
                    auxErro1 = fabs(XkD[tid]);

                tid+=n_threads;
            }

            __syncthreads();

            auxErro = auxErro / auxErro1;

            if(auxErro < erro)                         //Se o erro for menor que o erro que foi passado pelo usuario, acaba as iterações.
                break;

            tid = threadIdx.x;

            while(tid < ordem){
                XD[tid] = XkD[tid];
                tid+=n_threads;
            }

            tid = threadIdx.x;
            __syncthreads();
        }
        (*iteracaoD) = iteracao;
    }
}

int main(){
    int i, ordem, j, MaxIteracao, fila, iteracao, n_threads;

    scanf("%d", &ordem);
    scanf("%d", &fila);

    float *X, *Xk, erro;
    float *A, *B, *Ainicial;

    X = (float *) malloc(ordem*sizeof(float));
    Xk = (float *) malloc(ordem*sizeof(float));
    A = (float *) malloc(ordem*ordem*sizeof(float));
    Ainicial = (float *) malloc(ordem*ordem*sizeof(float));
    B = (float *) malloc(ordem*sizeof(float));

    int *iteracaoD;
    float *AD, *XD, *XkD;

    cudaMalloc((void**)&AD, ordem*ordem*sizeof(float));
    cudaMalloc((void**)&XD, ordem*sizeof(float));
    cudaMalloc((void**)&XkD, ordem*sizeof(float));
    cudaMalloc((void**)&iteracaoD, sizeof(int));

    scanf("%f", &erro);
    scanf("%d", &MaxIteracao);

    for(i=0; i<ordem*ordem; i++)  //Salva os valores da Matriz A.
        scanf("%f", &A[i]);

    for(i=0; i<ordem; i++)        //Salva os valores da Matriz B.
        scanf("%f", &B[i]);

    for(i=0; i<ordem*ordem; i++)  //Salva os valores da Matriz inicializa A, na auxiliar Matriz Ainicial.
        Ainicial[i] = A[i];

    for(i=0; i<ordem; i++)        //Guarda o valor da diagonal da Matriz A na matriz X.
        X[i] = A[i*ordem+i];

    for(i=0; i<ordem; i++)        //Calcula os valores para o calculo  das iterações.
        for(j=0; j<ordem; j++){
            if(i == j)
                A[i*ordem+i] = B[i] / X[i];
            else
                A[i*ordem+j] = -(A[i*ordem+j] / X[i]);
        }

    for(i=0; i<ordem; i++)       //Pega a matriz B (resultado) e dividi pela diagonal, para o calculo da iteraçao 0.
        X[i] = B[i] / X[i];

    cudaMemcpy(AD, A, ordem*ordem*sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(XD, X, ordem*sizeof(float), cudaMemcpyHostToDevice);

    n_threads = 1024;
    if(ordem < 1024)
        n_threads = ordem;

    dim3 grid(1,1,1);
    dim3 block(n_threads, 1, 1);

    jacobi<<<grid,block>>>(AD, XD, XkD, iteracaoD, ordem, erro, MaxIteracao, n_threads);

    cudaMemcpy(Xk, XkD, ordem*sizeof(float), cudaMemcpyDeviceToHost);
    cudaMemcpy(&iteracao, iteracaoD, sizeof(int), cudaMemcpyDeviceToHost);

    float Xkfinal = 0;         //Calculo do Xkfinal, que será o valor que resultará da conta com o valor proximo do X encontrado.
    for(i=0; i<ordem; i++)
        Xkfinal = Xkfinal + (Ainicial[fila*ordem+i] * Xk[i]);

    printf("Iterations: %d", iteracao);
    printf("\n");
    printf("RowTest: %d => [%f] =? %f", fila, Xkfinal, B[fila]);
    printf("\n");

    cudaFree(AD);
    cudaFree(XD);
    cudaFree(XkD);
    cudaFree(iteracaoD);
    free(X);
    free(Xk);
    free(B);
    free(A);
    free(Ainicial);

    return 0;
}
